import time

import torch
import torch.nn as nn
import torch.nn.functional as functional
import torchvision
import torchvision.transforms as transforms
from torch import GradScaler, autocast, optim


# 定义一个标准的CNN模型
class StaticCnn(nn.Module):
    def __init__(self):
        super(StaticCnn, self).__init__()
        self.conv1 = nn.Conv2d(3, 128, 7, stride=2, padding=3)
        self.bn1 = nn.BatchNorm2d(128)
        self.conv2 = nn.Conv2d(128, 256, 5, stride=1, padding=2)
        self.bn2 = nn.BatchNorm2d(256)
        self.conv3 = nn.Conv2d(256, 128, 3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 32, 5, 2, 2)
        self.bn4 = nn.BatchNorm2d(32)

        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = functional.relu(self.bn1(self.conv1(x)))
        x = functional.relu(self.bn2(self.conv2(x)))
        x = functional.relu(self.bn3(self.conv3(x)))
        x = functional.relu(self.bn4(self.conv4(x)))
        x = x.view(-1, 128 * 4 * 4)
        x = functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 训练模型的函数
def train_model(model_id, model, trainloader, criterion, optimizer, epochs=10):
    # 检查是否有可用的 GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 将模型转移到 GPU
    model.to(device)

    # 创建混合精度缩放器
    scaler = GradScaler()

    # 记录训练开始时间
    start_time = time.time()

    for epoch in range(epochs):
        model.train()  # 切换到训练模式
        running_loss = 0.0

        for inputs, labels in trainloader:
            # 将数据移到 GPU
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            # 混合精度前向传播
            with autocast(device_type='cuda'):
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            # 缩放后的反向传播
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()

        print(f'模型{model_id}: Epoch {epoch + 1}, Loss: {running_loss / len(trainloader):.4f}')

    # 计算训练时间
    end_time = time.time()
    training_time = end_time - start_time
    minutes = int(training_time // 60)  # 获取分钟数
    seconds = int(training_time % 60)  # 获取剩余的秒数

    print(f"模型{model_id}训练完毕, 耗时: {minutes} 分 {seconds} 秒")


# 测试模型的函数
def test_model(model_id, model, testloader):
    print("模型测试中...")
    print(model)

    # 确保模型在正确的设备上
    device = next(model.parameters()).device  # 获取模型所在的设备

    model.eval()  # 切换到评估模式

    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data

            # 将输入数据移到模型所在的设备
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            _, predicted = torch.max(outputs, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'模型{model_id}测试集准确率: {100 * correct / total:.2f}%')


def main():
    # 定义数据预处理方式（将图像转换为张量，并标准化）
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将图像转换为Tensor
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # 归一化
        transforms.RandomHorizontalFlip(),
    ])

    # 加载训练集和测试集
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)

    # 使用DataLoader来加载数据
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
    testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=4)

    # 获取类别名称
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    cnn = StaticCnn()

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()  # 分类问题的损失函数
    optimizer = optim.Adam(cnn.parameters(), lr=0.001)  # 使用Adam优化器
    train_model(1, cnn, trainloader, criterion, optimizer)
    test_model(1, cnn, testloader)


if __name__ == '__main__':
    main()
